//package org.jeecg.modules.airag.llm.tools;
//
//import dev.langchain4j.data.message.AiMessage;
//import dev.langchain4j.data.message.ChatMessage;
//import dev.langchain4j.data.message.SystemMessage;
//import dev.langchain4j.data.message.UserMessage;
//import dev.langchain4j.internal.Utils;
//import dev.langchain4j.internal.ValidationUtils;
//import dev.langchain4j.model.input.Prompt;
//import dev.langchain4j.model.input.PromptTemplate;
//import dev.langchain4j.rag.content.Content;
//import dev.langchain4j.rag.content.retriever.ContentRetriever;
//import dev.langchain4j.rag.query.Query;
//import net.sf.jsqlparser.JSQLParserException;
//import net.sf.jsqlparser.parser.CCJSqlParserUtil;
//import net.sf.jsqlparser.statement.select.Select;
//import org.slf4j.Logger;
//import org.slf4j.LoggerFactory;
//
//import javax.sql.DataSource;
//import java.sql.*;
//import java.util.*;
//
//public class MysqlDatabaseContentRetriever implements ContentRetriever {
//
//    private static final Logger log = LoggerFactory.getLogger(MysqlDatabaseContentRetriever.class);
//
//    private static final PromptTemplate TABLE_ANALYSIS_PROMPT = PromptTemplate.from(
//            "Available tables in the database:\n{{availableTables}}\n\nBased on the question, identify which tables from the above list are needed. Only output the table names separated by commas. If no relevant tables are found, output 'NONE'.");
//
//    private static final PromptTemplate SQL_GENERATION_PROMPT = PromptTemplate.from(
//            "Database structure:\n{{databaseStructure}}\nGenerate only SQL SELECT query for the question.today is  {{current_date}}.");
//
////    private final DataSource dataSource;
////    private final PromptTemplate promptTemplate;
////    private final ChatLanguageModel chatLanguageModel;
////    private final int maxRetries;
////    private static String catalog;
////    private Map<String, String> tableStructures;
////
////    /**
////     * 私有构造函数，用于创建MysqlDatabaseContentRetriever实例
////     *
////     * @param dataSource 数据源对象，用于建立数据库连接
////     * @param databaseStructure 数据库结构描述，如果为null则会自动生成
////     * @param promptTemplate 提示模板，用于生成SQL查询
////     * @param chatLanguageModel 聊天语言模型，用于生成SQL查询
////     * @param maxRetries 最大重试次数，默认为1
////     */
////    private MysqlDatabaseContentRetriever(DataSource dataSource, String databaseStructure,
////                                         PromptTemplate promptTemplate, ChatLanguageModel chatLanguageModel,
////                                         Integer maxRetries) {
////        this.dataSource = ValidationUtils.ensureNotNull(dataSource, "dataSource");
////        this.promptTemplate = Utils.getOrDefault(promptTemplate, SQL_GENERATION_PROMPT);
////        this.chatLanguageModel = ValidationUtils.ensureNotNull(chatLanguageModel, "chatLanguageModel");
////        this.maxRetries = Utils.getOrDefault(maxRetries, 1);
////        this.catalog = extractCatalogFromUrl(dataSource);
////        this.tableStructures = new HashMap<>();
////    }
//
//    /**
//     * 从数据源中提取数据库名称
//     *
//     * @param dataSource 数据源对象
//     * @return 数据库名称，如果无法获取则返回null
//     */
//    private static String extractCatalogFromUrl(DataSource dataSource) {
//        try (Connection connection = dataSource.getConnection()) {
//            System.out.println("connection.getCatalog() = " + connection.getCatalog());
//            return connection.getCatalog();
//        } catch (SQLException e) {
//            log.warn("无法获取数据库名称", e);
//            return null;
//        }
//    }
//
//    /**
//     * 生成数据库的DDL语句
//     * 遍历数据库中所有以'py'开头的表，生成它们的建表语句
//     *
//     * @param dataSource 数据源对象
//     * @return 包含所有表DDL语句的字符串
//     * @throws RuntimeException 当生成DDL过程中发生SQL异常时
//     */
//    private static String generateDDL(DataSource dataSource) {
//        log.info("开始生成数据库DDL");
//        StringBuilder ddl = new StringBuilder();
//        try (Connection connection = dataSource.getConnection()) {
//            DatabaseMetaData metaData = connection.getMetaData();
//            String catalog = connection.getCatalog();
//            ResultSet tables = metaData.getTables(catalog, null, "py%", new String[]{"TABLE"});
//            log.info("成功获取数据库连接和元数据信息");
//
//            while (tables.next()) {
//                String tableName = tables.getString("TABLE_NAME");
//                // 跳过系统表
//                if (tableName.startsWith("sys_") || tableName.startsWith("information_schema")) {
//                    log.info("跳过系统表: {}", tableName);
//                    continue;
//                }
//                log.info("正在处理表: {}", tableName);
//                String createTableStatement = generateCreateTableStatement(tableName, metaData);
//                ddl.append(createTableStatement).append("\n");
//            }
//        } catch (SQLException e) {
//            log.error("生成DDL时发生错误", e);
//            throw new RuntimeException(e);
//        }
//        log.info("DDL生成完成");
//        return ddl.toString();
//    }
//
//    /**
//     * 生成单个表的建表语句
//     * 包含列名、类型、是否可空、主键和注释等信息
//     *
//     * @param tableName 表名
//     * @param metaData 数据库元数据对象
//     * @return 表的建表语句
//     * @throws SQLException 当获取表结构信息过程中发生SQL异常时
//     */
//    private static String generateCreateTableStatement(String tableName, DatabaseMetaData metaData) throws SQLException {
//        log.info("开始生成表 {} 的建表语句", tableName);
//        try (Connection connection = metaData.getConnection();
//             Statement statement = connection.createStatement();
//             ResultSet resultSet = statement.executeQuery("SHOW CREATE TABLE " + tableName)) {
//
//            if (resultSet.next()) {
//                // 获取CREATE TABLE语句（在MySQL中是第二列）
//                String createTableStatement = resultSet.getString(2);
//                log.debug("成功获取表 {} 的建表语句", tableName);
//                log.info("成功生成表 {} 的建表语句", createTableStatement);
//                return createTableStatement + "\n";
//            } else {
//                log.warn("无法获取表 {} 的建表语句", tableName);
//                return "";
//            }
//        } catch (SQLException e) {
//            log.error("获取表 {} 的建表语句时发生错误", tableName, e);
//            throw e;
//        }
//    }
//
//    private static String generateCreateTableStatement1(String tableName, DatabaseMetaData metaData) throws SQLException {
//        log.info("开始生成表 {} 的建表语句", tableName);
//        StringBuilder createTableStatement = new StringBuilder();
//        ResultSet columns = metaData.getColumns(catalog, null, tableName, null);
//        ResultSet pk = metaData.getPrimaryKeys(catalog, null, tableName);
//
//        String primaryKeyColumn = "";
//        if (pk.next()) {
//            primaryKeyColumn = pk.getString("COLUMN_NAME");
//            log.debug("表 {} 的主键列为: {}", tableName, primaryKeyColumn);
//        }
//
//        createTableStatement.append("CREATE TABLE ").append(tableName).append(" (\n");
//
//        while (columns.next()) {
//            String columnName = columns.getString("COLUMN_NAME");
//            String columnType = columns.getString("TYPE_NAME");
//            String nullable = columns.getString("IS_NULLABLE").equals("YES") ? " NULL" : " NOT NULL";
//            String comment = columns.getString("REMARKS");
//
//            createTableStatement.append("  ").append(columnName)
//                    .append(" ").append(columnType)
//                    .append(nullable);
//
//            if (columnName.equals(primaryKeyColumn)) {
//                createTableStatement.append(" PRIMARY KEY");
//            }
//
//            if (comment != null && !comment.isEmpty()) {
//                createTableStatement.append(" COMMENT '").append(comment).append("'");
//            }
//
//            createTableStatement.append(",\n");
//        }
//
//        if (createTableStatement.charAt(createTableStatement.length() - 2) == ',') {
//            createTableStatement.delete(createTableStatement.length() - 2, createTableStatement.length());
//        }
//
//        createTableStatement.append("\n);\n");
//
//        return createTableStatement.toString();
//    }
//
//    /**
//     * 根据自然语言查询生成并执行SQL查询，返回查询结果
//     * 支持多次重试，每次重试会根据错误信息调整SQL查询
//     *
//     * @param query 自然语言查询对象
//     * @return 查询结果列表，如果查询失败则返回空列表
//     */
//    @Override
//    public List<Content> retrieve(Query query) {
//        log.info("开始处理查询: {}", query.text());
//        String sqlQuery = null;
//        String errorMessage = null;
//        int attemptsLeft = maxRetries + 1;
//        log.debug("剩余重试次数: {}", attemptsLeft);
//
//        while (attemptsLeft > 0) {
//            attemptsLeft--;
//            sqlQuery = generateSqlQuery(query, sqlQuery, errorMessage);
//            sqlQuery = clean(sqlQuery);
//
//            if (!isSelect(sqlQuery)) {
//                return Collections.emptyList();
//            }
//
//            try (Connection connection = dataSource.getConnection();
//                 Statement statement = connection.createStatement()) {
//                String result = execute(sqlQuery, statement);
//                log.info("成功执行SQL查询结果: {}", result);
//                Content content = format(result, sqlQuery);
//                return Collections.singletonList(content);
//            } catch (Exception e) {
//                errorMessage = e.getMessage();
//            }
//        }
//        return Collections.emptyList();
//    }
//
//    /**
//     * 生成SQL查询语句
//     * 使用聊天语言模型将自然语言转换为SQL查询
//     *
//     * @param query 自然语言查询对象
//     * @param previousSqlQuery 上一次失败的SQL查询，用于优化生成
//     * @param previousErrorMessage 上一次查询的错误信息
//     * @return 生成的SQL查询语句
//     */
//    private String generateSqlQuery(Query query, String previousSqlQuery, String previousErrorMessage) {
//        // 首先分析需要哪些表
//        List<String> requiredTables = analyzeRequiredTables(query.text());
//
//        // 获取所需表的DDL
//        String relevantDDL = generateRelevantDDL(requiredTables);
//
//        // 生成SQL查询
//        List<ChatMessage> messages = new ArrayList<>();
//        Map<String, Object> variables = new HashMap<>();
//        variables.put("databaseStructure", relevantDDL);
//        messages.add(SQL_GENERATION_PROMPT.apply(variables).toSystemMessage());
//        messages.add(UserMessage.from(query.text()));
//
//        if (previousSqlQuery != null && previousErrorMessage != null) {
//            messages.add(AiMessage.from(previousSqlQuery));
//            messages.add(UserMessage.from(previousErrorMessage));
//        }
////        return chatLanguageModel.chat(messages).aiMessage().text();
////        return messages;
//        return null;
//    }
//
//    /**
//     * 创建系统提示(确定需要的表名)
//     * 将数据库结构信息添加到提示模板中
//     *
//     * @return 包含数据库结构信息的提示对象
//     */
//    private List<String> analyzeRequiredTables(String question) {
//        try (Connection connection = dataSource.getConnection()) {
//            // 获取所有表名
//            DatabaseMetaData metaData = connection.getMetaData();
//            ResultSet tables = metaData.getTables(catalog, null, "py%", new String[]{"TABLE"});
//
//            // 构建可用表列表
//            StringBuilder availableTables = new StringBuilder();
//            while (tables.next()) {
//                String tableName = tables.getString("TABLE_NAME");
////                if (!tableName.startsWith("sys_") && !tableName.startsWith("information_schema")) {
////                    availableTables.append(tableName).append("\n");
////                }
//                availableTables.append(tableName).append("\n");
//            }
//
//            // 让模型选择相关表
//            List<ChatMessage> messages = new ArrayList<>();
//            Map<String, Object> variables = new HashMap<>();
//            variables.put("availableTables", availableTables.toString());
//            messages.add(TABLE_ANALYSIS_PROMPT.apply(variables).toSystemMessage());
//            messages.add(UserMessage.from(question));
////            String tableList = chatLanguageModel.chat(messages).aiMessage().text();
//
////            // 如果没有找到相关表，返回空列表
////            if ("NONE".equalsIgnoreCase(tableList.trim())) {
////                return Collections.emptyList();
////            }
////
////            return Arrays.asList(tableList.split(",\\s*"));
//            return null;
//        } catch (SQLException e) {
//            log.error("获取数据库表信息时发生错误", e);
//            return Collections.emptyList();
//        }
//    }
//
//    private String generateRelevantDDL(List<String> tableNames) {
//        StringBuilder ddl = new StringBuilder();
//        for (String tableName : tableNames) {
//            tableName = tableName.trim();
//            // 如果表结构已经缓存，直接使用
//            if (tableStructures.containsKey(tableName)) {
//                ddl.append(tableStructures.get(tableName)).append("\n");
//                continue;
//            }
//
//            try (Connection connection = dataSource.getConnection()) {
//                DatabaseMetaData metaData = connection.getMetaData();
//                String createTableStatement = generateCreateTableStatement(tableName, metaData);
//                tableStructures.put(tableName, createTableStatement);
//                ddl.append(createTableStatement).append("\n");
//            } catch (SQLException e) {
//                log.error("生成表 {} 的DDL时发生错误", tableName, e);
//            }
//        }
//        return ddl.toString();
//    }
//
//    /**
//     * 清理SQL查询字符串
//     * 移除Markdown代码块标记
//     *
//     * @param sqlQuery 原始SQL查询字符串
//     * @return 清理后的SQL查询字符串
//     */
//    private String clean(String sqlQuery) {
//        if (sqlQuery.contains("```sql")) {
//            return sqlQuery.substring(sqlQuery.indexOf("```sql") + 6, sqlQuery.lastIndexOf("```"));
//        } else if (sqlQuery.contains("```")) {
//            return sqlQuery.substring(sqlQuery.indexOf("```") + 3, sqlQuery.lastIndexOf("```"));
//        }
//        return sqlQuery;
//    }
//
//    /**
//     * 判断SQL查询是否为SELECT语句
//     *
//     * @param sqlQuery SQL查询字符串
//     * @return 如果是SELECT语句返回true，否则返回false
//     */
//    private boolean isSelect(String sqlQuery) {
//        try {
//            return CCJSqlParserUtil.parse(sqlQuery) instanceof Select;
//        } catch (JSQLParserException e) {
//            return false;
//        }
//    }
//
//    /**
//     * 执行SQL查询并格式化结果
//     * 将结果集转换为CSV格式字符串
//     *
//     * @param sqlQuery SQL查询语句
//     * @param statement SQL语句对象
//     * @return CSV格式的查询结果
//     * @throws SQLException 当执行SQL查询过程中发生异常时
//     */
//    private String execute(String sqlQuery, Statement statement) throws SQLException {
//        log.info("开始执行SQL查询: {}", sqlQuery);
//        List<String> resultRows = new ArrayList<>();
//        try (ResultSet resultSet = statement.executeQuery(sqlQuery)) {
//            int columnCount = resultSet.getMetaData().getColumnCount();
//            List<String> columnNames = new ArrayList<>();
//
//            for (int i = 1; i <= columnCount; i++) {
//                columnNames.add(resultSet.getMetaData().getColumnName(i));
//            }
//            resultRows.add(String.join(",", columnNames));
//
//            while (resultSet.next()) {
//                List<String> columnValues = new ArrayList<>();
//                for (int i = 1; i <= columnCount; i++) {
//                    String columnValue = resultSet.getObject(i) == null ? "" : resultSet.getObject(i).toString();
//                    if (columnValue.contains(",")) {
//                        columnValue = "\"" + columnValue + "\"";
//                    }
//                    columnValues.add(columnValue);
//                }
//                resultRows.add(String.join(",", columnValues));
//            }
//        }
//        return String.join("\n", resultRows);
//    }
//
//    /**
//     * 格式化查询结果
//     * 将SQL查询和结果组合成内容对象
//     *
//     * @param result 查询结果字符串
//     * @param sqlQuery 执行的SQL查询语句
//     * @return 格式化后的内容对象
//     */
//    private static Content format(String result, String sqlQuery) {
//        return Content.from(String.format("Result of executing '%s':\n%s", sqlQuery, result));
//    }
//
//    /**
//     * 创建构建器实例
//     * 用于构建MysqlDatabaseContentRetriever对象
//     *
//     * @return 新的构建器实例
//     */
//    public static MysqlDatabaseContentRetrieverBuilder builder() {
//        return new MysqlDatabaseContentRetrieverBuilder();
//    }
//
//    /**
//     * MysqlDatabaseContentRetriever的构建器类
//     * 用于以流式API方式构建MysqlDatabaseContentRetriever实例
//     */
//    public static class MysqlDatabaseContentRetrieverBuilder {
//        private DataSource dataSource;
//        private String databaseStructure;
//        private PromptTemplate promptTemplate;
//        private ChatLanguageModel chatLanguageModel;
//        private Integer maxRetries;
//
//        /**
//         * 设置数据源
//         *
//         * @param dataSource 用于数据库连接的数据源对象
//         * @return 构建器实例
//         */
//        public MysqlDatabaseContentRetrieverBuilder dataSource(DataSource dataSource) {
//            this.dataSource = dataSource;
//            return this;
//        }
//
//        /**
//         * 设置数据库结构描述
//         * 如果不设置，将自动从数据库生成
//         *
//         * @param databaseStructure 数据库结构的DDL描述
//         * @return 构建器实例
//         */
//        public MysqlDatabaseContentRetrieverBuilder databaseStructure(String databaseStructure) {
//            this.databaseStructure = databaseStructure;
//            return this;
//        }
//
//        /**
//         * 设置提示模板
//         * 用于生成SQL查询的提示模板
//         *
//         * @param promptTemplate 提示模板对象
//         * @return 构建器实例
//         */
//        public MysqlDatabaseContentRetrieverBuilder promptTemplate(PromptTemplate promptTemplate) {
//            this.promptTemplate = promptTemplate;
//            return this;
//        }
//
//        /**
//         * 设置聊天语言模型
//         * 用于将自然语言转换为SQL查询
//         *
//         * @param chatLanguageModel 聊天语言模型对象
//         * @return 构建器实例
//         */
//        public MysqlDatabaseContentRetrieverBuilder chatLanguageModel(ChatLanguageModel chatLanguageModel) {
//            this.chatLanguageModel = chatLanguageModel;
//            return this;
//        }
//
//        /**
//         * 设置最大重试次数
//         * 当SQL查询失败时的重试次数
//         *
//         * @param maxRetries 最大重试次数
//         * @return 构建器实例
//         */
//        public MysqlDatabaseContentRetrieverBuilder maxRetries(Integer maxRetries) {
//            this.maxRetries = maxRetries;
//            return this;
//        }
//
//        /**
//         * 构建MysqlDatabaseContentRetriever实例
//         *
//         * @return 新的MysqlDatabaseContentRetriever实例
//         */
//        public MysqlDatabaseContentRetriever build() {
//            return new MysqlDatabaseContentRetriever(dataSource, databaseStructure, promptTemplate,
//                    chatLanguageModel, maxRetries);
//        }
//    }
//}